draft for enable prefill in cudagraph #3354
Draft
+278
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Thanks for your contribution! |
Comment on lines
+1485
to
+1487
| if self.cudagraph_capture_prefill: | ||
| self.capture_model_prefill() | ||
|
|
There was a problem hiding this comment.
这个加在gpu_worker 里吧 和原本的 gpu_model_runner 是平级的
| @@ -1007,6 +1087,165 @@ def initialize_attn_backend(self) -> None: | |||
|
|
|||
| self.attn_backends.append(attn_backend) | |||
|
|
|||
| def _dummy_run_prefill( | |||
There was a problem hiding this comment.
整个dummy run都需要重写吗？是不是重写个_dummy_prefill_inputs_prefill就够了
| decode_exists = self.exist_decode() | ||
| paddle.distributed.all_gather_object(only_prefill_batch_list, not decode_exists) | ||
| only_prefill_batch = all(only_prefill_batch_list) | ||
| self.fd_config.parallel_config.moe_phase.phase = "prefill" if only_prefill_batch else "decode" |
There was a problem hiding this comment.
为什么需要改 moe_phase.phase
Comment on lines
+560
to
+573
| full_length = min( | ||
| num_tokens // batch_size, | ||
| self.parallel_config.max_model_len - max_dec_len, | ||
| ) | ||
|
|
||
| # NOTE(wanglongzhi): When the full length is too large, DeepEP's buffer size will not be enough to cause the result to appear nan. | ||
| # TODO(wanglongzhi): Figure out the accurate buffer size of DeepEP. | ||
| if self.fd_config.parallel_config.enable_expert_parallel: | ||
| full_length = min(full_length, 32) | ||
|
|
||
| input_length = int(full_length * self.cache_config.kv_cache_ratio) | ||
| block_num = ( | ||
| input_length + self.cache_config.block_size - 1 | ||
| ) // self.cache_config.block_size + self.cache_config.enc_dec_block_num |
There was a problem hiding this comment.
这部分逻辑能确保 input_length 等于想要捕获的 num_tokens 吗
| @@ -909,6 +967,28 @@ def initialize_forward_meta(self): | |||
| and not (prefill_exists if prefill_exists is not None else self.exist_prefill()) | |||
There was a problem hiding this comment.
seq_lens_encoder 这个tensor 指针会变，把get block shape kernel 输出的另外几个tensor 也打印出来看下
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.